package com.shenyaocn.android.Encoder;

import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaFormat;
import android.media.MediaMuxer;
import android.os.Build;

import java.io.IOException;
import java.nio.ByteBuffer;

/**
 * Created by ShenYao on 2014/6/21.
 */

public class AvcEncoder {

	private static final String MIME_TYPE = "video/avc";

	private MediaCodec.BufferInfo mVideoBufferInfo;
	private MediaCodec mEncoderVideo;
	private MediaMuxer mMuxer;
	private int mVideoTrackIndex = -1;

	private boolean mMuxerStarted = false;
	private String recordFileName;

	private byte[] yuvTempBuffer;
	private boolean isYuv420p = true;
	private boolean bHwUvReversed = false;

	private long startWhenVideo;
	private long lastEncodedVideoTimeStamp = 0;

	public static void NV21toI420(byte[] nv21, byte[] i420, int width, int height) {

		final int frameSize = width * height;
		final int qFrameSize = frameSize/4;

		System.arraycopy(nv21, 0, i420, 0, frameSize);

		for (int i = 0; i < qFrameSize; i++) {
			i420[frameSize + i] = nv21[frameSize + i*2 + 1];
			i420[frameSize + qFrameSize + i] = nv21[frameSize + i*2];
		}
	}

	public static void NV21toYV12(byte[] nv21, byte[] yv12, int width, int height) {

		final int frameSize = width * height;
		final int qFrameSize = frameSize/4;

		System.arraycopy(nv21, 0, yv12, 0, frameSize);

		for (int i = 0; i < qFrameSize; i++) {
			yv12[frameSize + qFrameSize + i] = nv21[frameSize + i*2 + 1];
			yv12[frameSize + i] = nv21[frameSize + i*2];
		}
	}

	public static void NV21toNV12(byte[] nv21, byte[] nv12, int width, int height){

		final int frameSize = width*height;
		int i = 0,j = 0;
		System.arraycopy(nv21, 0, nv12, 0, frameSize);

		for (j = 0; j < frameSize/2; j+=2) {
			nv12[frameSize + j + 1] = nv21[frameSize + j];
			nv12[frameSize + j] = nv21[frameSize + j + 1];
		}
	}

	public synchronized boolean open(String fileName, int width, int height) {
		if(Build.VERSION.SDK_INT < 18)
			return false;

		mMuxerStarted = false;

		try {
			prepareVideoEncoder(width, height);
			recordFileName = fileName;
			mMuxer = new MediaMuxer(fileName, MediaMuxer.OutputFormat.MUXER_OUTPUT_MPEG_4);
			return true;
		} catch (Exception ex) {
		}
		mMuxer = null;
		mEncoderVideo = null;

		return false;
	}

	public void setHwUvReversed(boolean bReversed) {
		bHwUvReversed = bReversed;
	}

	public String getRecordFileName() {
		return recordFileName;
	}

	private void prepareVideoEncoder(int width, int height) throws IOException {
		startWhenVideo = 0;
		mVideoBufferInfo = new MediaCodec.BufferInfo();

		MediaFormat format = MediaFormat.createVideoFormat(MIME_TYPE, width, height);

		format.setInteger(MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Planar);
		format.setInteger(MediaFormat.KEY_BIT_RATE, width * height * 6);
		format.setInteger(MediaFormat.KEY_FRAME_RATE, 30);
		format.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 1);

		mEncoderVideo = MediaCodec.createEncoderByType(MIME_TYPE);

		boolean bConf = false;
		try{
			mEncoderVideo.configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE); // 有些设备可能不支持，会引发异常
			bConf = true;
		} catch (Exception ex){}

		if(!bConf) {
			format.setInteger(MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar);
			mEncoderVideo.configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);
			isYuv420p = false;
		}

		mEncoderVideo.start();
	}

	public synchronized void close() {
		if(Build.VERSION.SDK_INT < 18)
			return;

		drainEncoder(mEncoderVideo, mVideoBufferInfo, mVideoTrackIndex, true);

		releaseEncoder();

		mMuxerStarted = false;
	}

	// 编码视频帧，里面包含yuv格式转换
	public synchronized void putFrame(byte[] pixelNv21, int width, int height) {
		if (mEncoderVideo == null) {
			return;
		}

		if(isYuv420p) {
			if (yuvTempBuffer == null || yuvTempBuffer.length != pixelNv21.length)
				yuvTempBuffer = new byte[pixelNv21.length];

			if(bHwUvReversed) {
				NV21toYV12(pixelNv21, yuvTempBuffer, width, height);
			} else {
				NV21toI420(pixelNv21, yuvTempBuffer, width, height);
			}

			VideoEncode(yuvTempBuffer, System.currentTimeMillis(), false);
		} else {

			if(!bHwUvReversed) {
				if (yuvTempBuffer == null || yuvTempBuffer.length != pixelNv21.length)
					yuvTempBuffer = new byte[pixelNv21.length];

				NV21toNV12(pixelNv21, yuvTempBuffer, width, height);
				VideoEncode(yuvTempBuffer, System.currentTimeMillis(), false);
			} else {
				VideoEncode(pixelNv21, System.currentTimeMillis(), false);
			}
		}
		mVideoTrackIndex = drainEncoder(mEncoderVideo, mVideoBufferInfo, mVideoTrackIndex, false);
	}

	public boolean isOpened() {
		return !(mMuxer == null);
	}

	public boolean isStarted() {
		return mMuxerStarted;
	}

	private void releaseEncoder() {
		if (mEncoderVideo != null) {
			mEncoderVideo.stop();
			mEncoderVideo.release();
			mEncoderVideo = null;
		}

		if (mMuxer != null) {
			mMuxer.stop();
			mMuxer.release();
			mMuxer = null;
		}
	}

	private void VideoEncode(final byte[] pixelsNv21, long presentationTimeMs, boolean endOfStream) {

		if (mEncoderVideo == null) {
			return;
		}

		try{
			final ByteBuffer[] inputBuffers = mEncoderVideo.getInputBuffers();
			final int inputBufferIndex = mEncoderVideo.dequeueInputBuffer(10000);
			if (inputBufferIndex >= 0) {
				final ByteBuffer inputBuffer = inputBuffers[inputBufferIndex];
				inputBuffer.clear();
				inputBuffer.put(pixelsNv21);

				long duration = 0;

				if(startWhenVideo == 0) {
					startWhenVideo = System.currentTimeMillis();
				} else {
					duration = (presentationTimeMs - startWhenVideo) * 1001;
				}

				if (endOfStream) {
					mEncoderVideo.queueInputBuffer(inputBufferIndex, 0, pixelsNv21.length, duration, MediaCodec.BUFFER_FLAG_END_OF_STREAM);
				} else {
					mEncoderVideo.queueInputBuffer(inputBufferIndex, 0, pixelsNv21.length, duration, 0);
				}
			}
		} catch (Exception ex){

		}
	}

	private int drainEncoder(MediaCodec mediaCodec, MediaCodec.BufferInfo bufferInfo, int trackIndex, boolean endOfStream) {
		if (mMuxer == null || mediaCodec == null) {
			return trackIndex;
		}

		if (endOfStream) {
			if(mediaCodec == mEncoderVideo) {
				//mediaCodec.signalEndOfInputStream();
				VideoEncode(new byte[0], System.currentTimeMillis(), true);
			}
		}

		final int TIMEOUT_USEC = 1000; // 设定超时，一般情况下用不到，可以设置成-1让编码器无限等待，不过要是遇上一些奇葩的设备可能app就卡死了
		try {
			ByteBuffer[] encoderOutputBuffers = mediaCodec.getOutputBuffers();
			while (true) {
				int encoderStatus = mediaCodec.dequeueOutputBuffer(bufferInfo, TIMEOUT_USEC);
				if (encoderStatus == MediaCodec.INFO_TRY_AGAIN_LATER) {

					if (!endOfStream) {
						break;
					}
				} else if (encoderStatus == MediaCodec.INFO_OUTPUT_BUFFERS_CHANGED) {

					encoderOutputBuffers = mediaCodec.getOutputBuffers();
				} else if (encoderStatus == MediaCodec.INFO_OUTPUT_FORMAT_CHANGED) {
					if (mMuxerStarted) {
						throw new RuntimeException("format changed twice");
					}

					if(mediaCodec == mEncoderVideo) {
						MediaFormat newFormat = mediaCodec.getOutputFormat();
						mVideoTrackIndex = mMuxer.addTrack(newFormat);
						trackIndex = mVideoTrackIndex;
						mMuxer.start();
						lastEncodedVideoTimeStamp = 0;
						mMuxerStarted = true;
					}
				} else if (encoderStatus < 0) {

				} else {
					ByteBuffer encodedData = encoderOutputBuffers[encoderStatus];
					if (encodedData == null) {
						throw new RuntimeException("encoderOutputBuffer " + encoderStatus +
								" was null");
					}

					if ((bufferInfo.flags & MediaCodec.BUFFER_FLAG_CODEC_CONFIG) != 0) {
						bufferInfo.size = 0;
					}

					if (bufferInfo.size != 0) {
						if (!mMuxerStarted) {
							throw new RuntimeException("muxer hasn't started");
						}

						encodedData.position(bufferInfo.offset);
						encodedData.limit(bufferInfo.offset + bufferInfo.size);

						if(mediaCodec == mEncoderVideo){
							if(bufferInfo.presentationTimeUs < lastEncodedVideoTimeStamp)
								bufferInfo.presentationTimeUs = lastEncodedVideoTimeStamp += 100000;
							lastEncodedVideoTimeStamp = bufferInfo.presentationTimeUs;
						}
						if(bufferInfo.presentationTimeUs < 0){
							bufferInfo.presentationTimeUs = 0;
						}
						mMuxer.writeSampleData(trackIndex, encodedData, bufferInfo);
					}

					mediaCodec.releaseOutputBuffer(encoderStatus, false);

					if ((bufferInfo.flags & MediaCodec.BUFFER_FLAG_END_OF_STREAM) != 0) {
						break;
					}
				}
			}
		} catch (Exception ex) {
		}
		return trackIndex;
	}
}